import numpy as np
import cv2
import glob

#设置终止条件，迭代30次或变动小于0.001
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
#生成42x3的矩阵，用来保存棋盘图中6*7个内焦点的3D坐标，也就是物体的点坐标
objp = np.zeros((6*7, 3), np.float32)
#通过np.mgrid生成对象的xy坐标点
#最终得到的objp为（0,0,0），（1,0,0），（2,0,0），。。。（6,5,0）
objp[:, :2] = np.mgrid[0:7, 0:6].T.reshape(-1, 2)
#使用CV函数求图像中的角点坐标，然后用CornerSubPixiv进行优化
#用于保存物体点
obj_points = []
#用于保存图像点
img_points = []
#返回当前目录下所有匹配的jpg图像
images = glob.glob('*.jpg')
for fname in images:
    #读取图像
    img = cv2.imread(fname)
    #转换为灰度图像
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    #寻找棋盘图的内角点位置
    ret, corners = cv2.findChessboardCorners(gray, (7, 6), None)
    #如果找到的棋盘图内所有内角点
    if ret == True:
        obj_points.append(objp)
        #亚像素级角点检测，在角点检测中精确化角点的位置
        corners2 = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        img_points.append(corners2)
        #在图像中标注角点，方便查看
        img = cv2.drawChessboardCorners(img, (7, 6), corners, ret)
        # cv2.imshow('img', img)
        # cv2.waitKey(0)
cv2.destroyAllWindows()
#下面使用cv2.calibrateCamera()进行相机标定，返回相机矩阵，畸变系数，旋转和平移向量
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(obj_points, img_points,
                                                   gray.shape, None, None)
# print(ret, '\n', mtx, '\n', dist, '\n', rvecs, '\n', tvecs)
'''
0.1553690691007881 

 [[534.07088364   0.         341.53407554]
 [  0.         534.11914595 232.94565259]
 [  0.           0.           1.        ]] 

 [[-2.92971637e-01  1.07706962e-01  1.31038376e-03 -3.11018780e-05
   4.34798110e-02]] 
'''
#消除畸变
#读取图像
img = cv2.imread('left12.jpg')
#获取图像的长宽
h, w = img.shape[:2]
#根据尺度因子调节相机矩阵
newcameratx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))
#在OpenCV中消除畸变的方法有两种
#方法一：调用undistort函数
dst1 = cv2.undistort(img, mtx, dist, None, newcameratx)
#裁剪图像
x, y, w, h = roi
dst1 = dst1[y:y+h, x:x+w]
cv2.imshow('dst1', dst1)
# cv2.waitKey(0)
#方法二：调用remap函数
#矫正畸变图像
mapx, mapy = cv2.initUndistortRectifyMap(mtx, dist, None, newcameratx, (w, h), 5)
dst2 = cv2.remap(img, mapx, mapy, cv2.INTER_LINEAR)
#截取图像
x, y, w, h = roi
dst2 = dst2[y:y+h, x:x+w]
cv2.imshow('dst2', dst2)
cv2.waitKey(0)
cv2.destroyAllWindows()
#计算重投影
mean_error = 0
for i in range(len(obj_points)):
    img_points2, _ = cv2.projectPoints(obj_points[i], rvecs[i],
                                       tvecs[i], mtx, dist)
    error = cv2.norm(img_points[i], img_points2, cv2.NORM_L2)/len(img_points2)
    mean_error += error
print("total error：", mean_error/len(obj_points))